import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from mpl_toolkits.mplot3d import Axes3D

# -------------------------- 基础设置 --------------------------
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False

# -------------------------- 1. 物理参数定义 --------------------------
target_true = np.array([0, 200, 0])    # 真目标下底面圆心
r, h = 7, 10                           # 真目标尺寸：半径7m、高10m
target_fake = np.array([0, 0, 0])      # 假目标位置

missile_M1 = np.array([20000, 0, 2000])# 导弹初始位置
missile_speed = 300                    # 导弹速度(m/s)
drone_FY1 = np.array([17800, 0, 1800]) # 无人机初始位置

g = 9.8                                # 重力加速度(m/s²)
smoke_radius = 10                      # 烟幕有效半径(m)
smoke_sink_speed = 3                   # 烟幕下沉速度(m/s)
smoke_valid_time = 20                  # 烟幕有效时间(s)

# 可优化参数（遗传算法待优化变量）
param_bounds = {
    'drone_speed': (100, 150),         # 无人机速度范围(m/s)
    'drop_delay': (1.0, 3.0),          # 投放延迟范围(s)
    'detonate_delay': (2.0, 5.0)       # 起爆延迟范围(s)
}


# -------------------------- 2. 核心计算函数 --------------------------
def get_cylinder_bounding_box_vertices(target_center, radius, height):
    """生成真目标外接长方体顶点"""
    x_min, x_max = target_center[0] - radius, target_center[0] + radius
    y_min, y_max = target_center[1] - radius, target_center[1] + radius
    z_min, z_max = target_center[2], target_center[2] + height
    return [
        np.array([x_min, y_min, z_min]), np.array([x_max, y_min, z_min]),
        np.array([x_max, y_max, z_min]), np.array([x_min, y_max, z_min]),
        np.array([x_min, y_min, z_max]), np.array([x_max, y_min, z_max]),
        np.array([x_max, y_max, z_max]), np.array([x_min, y_max, z_max])
    ]


def missile_real_time_position(missile_init_pos, t):
    """计算t时刻导弹位置"""
    if t < 0:
        return missile_init_pos
    dir_vec = target_fake - missile_init_pos
    dir_vec_mag = np.linalg.norm(dir_vec)
    dir_unit_vec = dir_vec / dir_vec_mag if dir_vec_mag > 1e-6 else np.array([0, 0, 0])
    return missile_init_pos + missile_speed * t * dir_unit_vec


def get_smoke_parameters(drone_speed, drop_delay, detonate_delay):
    """计算烟幕关键参数（起爆时刻、位置）"""
    detonate_t = drop_delay + detonate_delay  # 起爆时刻
    # 投放点坐标
    drop_x = drone_FY1[0] - drone_speed * drop_delay
    # 起爆点坐标（平抛运动）
    explosion_x = drop_x - drone_speed * detonate_delay
    fall_distance = 0.5 * g * (detonate_delay **2)
    explosion_z = max(drone_FY1[2] - fall_distance, 0.0)
    explosion_center = np.array([explosion_x, drone_FY1[1], explosion_z])
    return detonate_t, explosion_center


def smoke_real_time_center(t, detonate_t, explosion_center):
    """计算t时刻烟幕中心位置"""
    if t < detonate_t:
        return np.array([0, 0, 0])
    sink_distance = smoke_sink_speed * (t - detonate_t)
    real_z = max(explosion_center[2] - sink_distance, 0.0)
    return np.array([explosion_center[0], explosion_center[1], real_z])


def is_segment_sphere_intersect(missile_pos, vertex_pos, smoke_center, radius):
    """判断线段与球体是否相交"""
    vec_AC = smoke_center - missile_pos
    vec_AB = vertex_pos - missile_pos
    ab_sq_mag = np.dot(vec_AB, vec_AB)
    if ab_sq_mag < 1e-12:
        return False
    k = np.dot(vec_AC, vec_AB) / ab_sq_mag
    
    if k <= 0.0:
        min_dist = np.linalg.norm(vec_AC)
    elif 0.0 < k < 1.0:
        cross_prod = np.cross(vec_AC, vec_AB)
        min_dist = np.linalg.norm(cross_prod) / np.sqrt(ab_sq_mag)
    else:
        min_dist = np.linalg.norm(smoke_center - vertex_pos)
    return min_dist <= radius + 1e-6


def calculate_occlusion_duration(drone_speed, drop_delay, detonate_delay, time_step=0.01):
    """计算给定参数下的有效遮蔽时长"""
    # 获取烟幕参数
    detonate_t, explosion_center = get_smoke_parameters(drone_speed, drop_delay, detonate_delay)
    smoke_start = detonate_t
    smoke_end = detonate_t + smoke_valid_time
    target_vertices = get_cylinder_bounding_box_vertices(target_true, r, h)
    
    total_duration = 0.0
    prev_effective = False
    start_t = 0.0
    
    t = smoke_start
    while t <= smoke_end + 1e-6:
        # 计算当前状态
        missile_pos = missile_real_time_position(missile_M1, t)
        smoke_center = smoke_real_time_center(t, detonate_t, explosion_center)
        current_effective = any([
            is_segment_sphere_intersect(missile_pos, v, smoke_center, smoke_radius)
            for v in target_vertices
        ])
        
        # 状态变化处理
        if current_effective and not prev_effective:
            start_t = t
        elif not current_effective and prev_effective:
            total_duration += t - start_t
        prev_effective = current_effective
        t += time_step
    
    if prev_effective:
        total_duration += smoke_end - start_t
    return max(total_duration, 0.0)  # 确保非负


# -------------------------- 3. 数据生成与拟合（为遗传算法准备） --------------------------
def generate_training_data(num_samples=100):
    """生成参数-结果样本数据"""
    np.random.seed(42)  # 固定种子确保可复现
    data = []
    
    # 随机生成参数组合
    for _ in range(num_samples):
        drone_speed = np.random.uniform(*param_bounds['drone_speed'])
        drop_delay = np.random.uniform(*param_bounds['drop_delay'])
        detonate_delay = np.random.uniform(*param_bounds['detonate_delay'])
        
        # 计算对应遮蔽时长
        duration = calculate_occlusion_duration(drone_speed, drop_delay, detonate_delay)
        data.append([drone_speed, drop_delay, detonate_delay, duration])
    
    return np.array(data)


def fit_duration_model(X, y, degree=3):
    """拟合参数到遮蔽时长的多项式模型"""
    # 生成多项式特征（参数组合的交叉项）
    from sklearn.preprocessing import PolynomialFeatures
    poly = PolynomialFeatures(degree=degree, include_bias=False)
    X_poly = poly.fit_transform(X)
    
    # 线性回归拟合
    from sklearn.linear_model import LinearRegression
    model = LinearRegression()
    model.fit(X_poly, y)
    
    # 评估拟合效果
    y_pred = model.predict(X_poly)
    mse = np.mean((y - y_pred)** 2)
    print(f"拟合均方误差(MSE): {mse:.6f}")
    print(f"拟合决定系数(R²): {model.score(X_poly, y):.6f}")
    
    return model, poly


def predict_duration(model, poly, params):
    """使用拟合模型预测遮蔽时长"""
    params = np.array(params).reshape(1, -1)
    params_poly = poly.transform(params)
    return model.predict(params_poly)[0]


# -------------------------- 4. 遗传算法接口（预留） --------------------------
def genetic_algorithm_fitness(params):
    """遗传算法适应度函数（调用拟合模型快速计算）"""
    # params格式: [drone_speed, drop_delay, detonate_delay]
    # 确保参数在有效范围内
    for i, key in enumerate(param_bounds.keys()):
        params[i] = np.clip(params[i], *param_bounds[key])
    
    # 使用拟合模型预测（第二题可直接替换为真实计算）
    return predict_duration(fitted_model, poly_transformer, params)


# -------------------------- 5. 执行与可视化 --------------------------
if __name__ == "__main__":
    # 生成训练数据（可根据需要增加样本量）
    print("正在生成训练数据...")
    training_data = generate_training_data(num_samples=200)
    X = training_data[:, :3]  # 参数：无人机速度、投放延迟、起爆延迟
    y = training_data[:, 3]   # 目标：遮蔽时长
    
    # 拟合模型
    print("\n正在拟合参数-时长模型...")
    fitted_model, poly_transformer = fit_duration_model(X, y, degree=3)
    
    # 可视化拟合效果
    y_pred = [predict_duration(fitted_model, poly_transformer, x) for x in X]
    plt.figure(figsize=(10, 6))
    plt.scatter(y, y_pred, alpha=0.6)
    plt.plot([y.min(), y.max()], [y.min(), y.max()], 'r--')
    plt.xlabel('真实遮蔽时长(s)')
    plt.ylabel('拟合预测时长(s)')
    plt.title('拟合模型效果验证')
    plt.grid(True, alpha=0.3)
    plt.show()
    
    # 保存模型（为第二题做准备）
    import joblib
    joblib.dump((fitted_model, poly_transformer, param_bounds), 'occlusion_model.pkl')
    print("\n模型已保存为 'occlusion_model.pkl'，可用于第二题遗传算法")
